
#include "PsFrameDemuxer.h"
//#include "ConLog.h"

#include <stdio.h>
#include <memory.h>



#define PAYLOAD_TYPE_H264              (98)
#define PAYLOAD_TYPE_H265              (120)
#define PAYLOAD_TYPE_G711A             (8)
#define PAYLOAD_TYPE_G711U             (0)


#define PS_STREAM_TYPE_VIDEO_MPEG1     0x01
#define PS_STREAM_TYPE_VIDEO_MPEG2     0x02
#define PS_STREAM_TYPE_AUDIO_MPEG1     0x03
#define PS_STREAM_TYPE_AUDIO_MPEG2     0x04
#define PS_STREAM_TYPE_PRIVATE_SECTION 0x05
#define PS_STREAM_TYPE_PRIVATE_DATA    0x06
#define PS_STREAM_TYPE_AUDIO_AAC       0x0f
#define PS_STREAM_TYPE_AUDIO_AAC_LATM  0x11
#define PS_STREAM_TYPE_VIDEO_MPEG4     0x10
#define PS_STREAM_TYPE_VIDEO_H264      0x1b
#define PS_STREAM_TYPE_VIDEO_HEVC      0x24
#define PS_STREAM_TYPE_VIDEO_CAVS      0x42
#define PS_STREAM_TYPE_VIDEO_VC1       0xea
#define PS_STREAM_TYPE_VIDEO_DIRAC     0xd1
#define PS_STREAM_TYPE_AUDIO_AC3       0x81
#define PS_STREAM_TYPE_AUDIO_DTS       0x82
#define PS_STREAM_TYPE_AUDIO_TRUEHD    0x83
#define PS_STREAM_TYPE_AUDIO_G711A     0x90
#define PS_STREAM_TYPE_AUDIO_G711U     0x91

#define GETU16LEN(buf)      (((buf)[0] << 8) | (buf)[1])

static int gFileSize = 0;

// --------------------------------------------------------------------
// PSFrameDemuxer

PsFrameDemuxer::PsFrameDemuxer(PSFrameDemuxerCallback * pCallback)
{
    m_pCallback = pCallback;

    memset(&m_sPsContext, 0, sizeof(PsContext));
}

PsFrameDemuxer::~PsFrameDemuxer()
{

}

void PsFrameDemuxer::GetStatisticCnt(unsigned & uPsPacketCnt, unsigned & uVideoPacketCnt, unsigned & uAudioPacketCnt)
{
    uPsPacketCnt = m_uPsPacketCnt;
    uVideoPacketCnt = m_uVideoPacketCnt;
    uAudioPacketCnt = m_uAudioPacketCnt;
}

void PsFrameDemuxer::GetMediaStreamType(int & iVideoStreamType, int & iAudioStreamType)
{
    iVideoStreamType = m_iVideoStreamType;
    iAudioStreamType = m_iAudioStreamType;
}

void PsFrameDemuxer::ResetStatisticParam()
{
    m_uPsPacketCnt = 0;
    m_uVideoPacketCnt = 0;
    m_uAudioPacketCnt = 0;

    m_iVideoStreamType = -1;
    m_iAudioStreamType = -1;
}

void PsFrameDemuxer::OnPsFrame(const unsigned char * pszBuffer, unsigned bufferSize, unsigned long timestamp)
{
    if (NULL == pszBuffer || bufferSize < 1)
    {
        return;
    }

    m_uPsPacketCnt++;

    int parsed = 0;

    while (1)
    {
        int size = ParsePsFrame(pszBuffer + parsed, bufferSize - parsed, timestamp);
        if (size <= 0)
        {
            break;
        }

        parsed += size;
    }
}


int PsFrameDemuxer::ParsePsFrame(const unsigned char * pszBuffer, unsigned bufferSize, unsigned long timestamp)
{
    int size = PSFrameDemux(&m_sPsContext, pszBuffer, bufferSize);

    if (size <= 0 || !m_sPsContext.is_pes)
    {
        return size;
    }

    for (int i = 0; i < m_sPsContext.info.prog[0].stream_num; ++i)
    {
        int streamNo = m_sPsContext.stream_no;
        unsigned char streamId = m_sPsContext.info.prog[0].stream[i].stream_id;
        
        unsigned char streamType = m_sPsContext.info.prog[0].stream[i].stream_type;

        if (0xE0 == streamId && streamNo == i) // Video
        {
            m_uVideoPacketCnt++;

            m_iVideoStreamType = streamType;
            
            if ((PS_STREAM_TYPE_VIDEO_H264 == streamType) || (PS_STREAM_TYPE_VIDEO_HEVC == streamType))
            {
                gFileSize += m_sPsContext.es_len;
                if (m_pCallback)
                {
                    m_pCallback->OnVideoFrame(m_sPsContext.es_ptr, m_sPsContext.es_len, streamType, timestamp);
                }
            }
            else
            {
                printf("PsFrameDemuxer::ParsePsFrame (STREAM_TYPE_VIDEO_H264 != streamType && STREAM_TYPE_VIDEO_HEVC != streamType)");
            }
        }
        else if (0xC0 == streamId && streamNo == i) // audio
        {
            m_uAudioPacketCnt++;

            m_iAudioStreamType = streamType;
            
            if ((PS_STREAM_TYPE_AUDIO_G711A == streamType) || (PS_STREAM_TYPE_AUDIO_G711U == streamType))
            {
                if (m_pCallback)
                {
                    m_pCallback->OnAudioFrame(m_sPsContext.es_ptr, m_sPsContext.es_len, streamType, timestamp);
                }
            }
            else
            {
                printf("PsFrameDemuxer::ParsePsFrame (STREAM_TYPE_AUDIO_G711A != streamType && STREAM_TYPE_AUDIO_G711U != streamType)");
                if (m_pCallback)
				{
					m_pCallback->OnAudioFrame(m_sPsContext.es_ptr, m_sPsContext.es_len, streamType, timestamp);
				}
            }
        }
    }

    return size;
}

int PsFrameDemuxer::PSFrameDemux(PsContext * context, const unsigned char * pszBuffer, unsigned bufferSize)
{
    if (!pszBuffer || bufferSize < 1)
    {
        return -1;
    }

    int ret_len = 0;

    m_sPsContext.is_pes = 0;
    m_sPsContext.program_no = 0;
    m_sPsContext.stream_no = 0;
    m_sPsContext.pts = 0;
    m_sPsContext.pack_ptr = NULL;
    m_sPsContext.pack_len = 0;
    m_sPsContext.pes_head_len = 0;
    m_sPsContext.es_ptr = NULL;
    m_sPsContext.es_len = 0;

    int i = 0;
    for (i = 0; i < (int)bufferSize - 4; ++i)
    {
        ret_len = i;

        if (pszBuffer[i] == 0 && pszBuffer[i + 1] == 0 && pszBuffer[i + 2] == 1)
        {
            int type = pszBuffer[i + 3];

            if (type < 0xB9) // PS stream_id必须大于0xB9
            {
                continue;
            }

            if (!m_sPsContext.ps_started)// 未找到ps头
            {
                if (type != 0xBA) // 不是ps头
                {
                    continue; // 略过，继续往下搜寻
                }

                m_sPsContext.ps_started = 1; // 找到ps头
            }

            unsigned char * tmpBuffer = const_cast<unsigned char *>(pszBuffer);

            switch (type) //
            {
            case 0xBA: // PS header
                ret_len = HandlePsHeader(&m_sPsContext, tmpBuffer + i, bufferSize - i);
                break;
            case 0xBB: // PS system header
                ret_len = HandlePsSystemHeader(&m_sPsContext, tmpBuffer + i, bufferSize - i);
                break;
            case 0xBC: // PS map
                ret_len = HandlePsSystemMap(&m_sPsContext, tmpBuffer + i, bufferSize - i);
                break;
            case 0xB9: // PS finish
                ret_len = HandlePsFinish(&m_sPsContext, tmpBuffer + i, bufferSize - i);
                break;
            default:   // PES
                ret_len = HandlePes(&m_sPsContext, tmpBuffer + i, bufferSize - i);
                break;
            }

            break; //
        }
    }

    if (!m_sPsContext.ps_started)
    {
        return -1;
    }

    if (i + ret_len > (int)bufferSize) // 如果超过输入长度，返回0，下次再解析
    {
        return 0;
    }

    m_sPsContext.pack_ptr = const_cast<unsigned char *>(pszBuffer)+i;
    m_sPsContext.pack_len = ret_len;

    return i + ret_len;
}

int PsFrameDemuxer::HandlePsHeader(PsContext * context, unsigned char * buffer, int bufSize)
{
    if (bufSize < int(sizeof(PsPackHeader)))
    {
        return 0;
    }

    PsPackHeader * head = (PsPackHeader *)buffer;

    int staffingLen = head->pack_stuffing_length & 0x07;
    int psize = sizeof(PsPackHeader) + staffingLen;

    if (bufSize < psize)
    {
        return 0;
    }

    return psize;
}

int PsFrameDemuxer::HandlePsSystemHeader(PsContext * context, unsigned char * buffer, int bufSize)
{
    int psize = SizeCheck(buffer, bufSize);

    return psize;
}

int PsFrameDemuxer::HandlePsSystemMap(PsContext * context, unsigned char * buffer, int bufSize)
{
    int psize = SizeCheck(buffer, bufSize);

    if (psize <= 0)
    {
        return psize;
    }

    PsSystemMap * head = (PsSystemMap *)buffer;

    int info_len = GETU16LEN(head->ps_info_length);
    int map_len = GETU16LEN(head->es_map_length + info_len);

    unsigned char * map_buf = buffer + sizeof(PsSystemMap) + info_len;

    int pos = 0;
    int i = 0;

    context->info.program_num = 1;

    while (pos < map_len)
    {
        PsSystemMapEs * mes = (PsSystemMapEs *)(map_buf + pos);
        int es_info_len = GETU16LEN(mes->es_info_length);
        if (es_info_len < 0)
        {
            break;
        }

        if (i >= MAX_STREAM_NUM)
        {
            break;
        }


        pos += (sizeof(PsSystemMapEs) + es_info_len);

        // 如果属于音频或视频


        if ((mes->es_id & 0xC0) == 0xC0 || (mes->es_id & 0xE0) == 0xE0)
        {
            context->info.prog[0].stream[i].stream_type = mes->stream_type;
            context->info.prog[0].stream[i].stream_id = mes->es_id;

            i++;

            //printf("map stream_id=0x%x, stream_type=0x%x\n", mes->stream_type, mes->es_id);
        }
    }

    context->info.prog[0].stream_num = i;

    return psize;
}

int PsFrameDemuxer::HandlePsFinish(PsContext * context, unsigned char * buffer, int bufSize)
{
    int psize = 4;

    return psize;
}

int PsFrameDemuxer::HandlePes(PsContext * context, unsigned char * buffer, int bufSize)
{
    int psize = SizeCheck(buffer, bufSize);
    if (psize <= 0)
    {
        return psize;
    }

    unsigned char stream_id;
    unsigned long long pts;
    int es_len;

    int pes_head_len = HandlePesHeader(buffer, bufSize, &stream_id, &pts, &es_len);
    if (pes_head_len <= 0)
    {
        return psize;

    }

    if (context->info.program_num == 1)
    {
        for (int i = 0; i < context->info.prog[0].stream_num; i++)
        {
            if (stream_id == context->info.prog[0].stream[i].stream_id)
            {

                context->is_pes = 1;
                context->program_no = 0;
                context->stream_no = i;
                context->pts = pts;
                context->pes_head_len = pes_head_len;
                context->es_ptr = buffer + pes_head_len;
                context->es_len = es_len;

                const char startcode[4] = { 0, 0, 0, 1 };
                if (es_len != 320 && memcmp(startcode, context->es_ptr, 4))
                {
                    int debugx = 1;
                }

                break;
            }
        }
    }

    return psize;
}

int PsFrameDemuxer::HandlePesHeader(unsigned char * pesBuffer, int pesBufSize, unsigned char * streamId, unsigned long long * pts, int * esSize)
{
    int pesHeadSize = GetPesHeadSize(pesBuffer, pesBufSize);
    if (pesHeadSize <= 0)
    {
        return 0;
    }

    if (streamId)
    {
        *streamId = pesBuffer[3];
    }

    if (pts) // 解析PTS
    {
        unsigned char flags_2 = pesBuffer[7];
        if (flags_2 & 0x80)
        {
            unsigned char * pts_buf = &pesBuffer[9];

            *pts = ((unsigned long long)pts_buf[0] & 0x0E) << 29;
            *pts |= ((unsigned long long)pts_buf[1]) << 22;
            *pts |= ((unsigned long long)pts_buf[2] & 0xFE) << 14;
            *pts |= ((unsigned long long)pts_buf[3]) << 7;
            *pts |= ((unsigned long long)pts_buf[4] & 0xFE) >> 1;
            *pts /= 90;
        }
    }

    if (esSize) // 解析ES长度
    {
        int pes_len = GETU16LEN(&pesBuffer[4]);

        int pes_head_len = pesBuffer[8];

        *esSize = pes_len - 3 - pes_head_len;

        //printf("pes_len=0x%x,pes_head_len=0x%x,es_len=0x%x\n",pes_len,pes_head_len,*es_len);
    }

    return pesHeadSize;
}

int PsFrameDemuxer::GetPesHeadSize(unsigned char * pesBuffer, int pesBufSize)
{
    int pesHeadSize = 0;

    if (pesBufSize < 9)
    {
        return 0;
    }

    if (pesBuffer[0] == 0 && pesBuffer[1] == 0 && pesBuffer[2] == 1)
    {
        if ((pesBuffer[3] & 0xC0) || (pesBuffer[3] & 0xE0))
        {
            pesHeadSize = 9 + pesBuffer[8];
        }
    }

    return pesHeadSize;
}

int PsFrameDemuxer::SizeCheck(unsigned char * buffer, int bufSize)
{
    int psize = 0;

    if (bufSize < 6)
    {
        return 0;
    }

    psize = 6 + GETU16LEN(&buffer[4]);

    if (bufSize < psize)
    {
        return 0;
    }

    return psize;
}
